@echo off
echo ================================
echo 🚀 本地 Whisper GPU 加速安装
echo ================================
echo.

REM 检查 Python 是否可用
echo 🔍 检查 Python 环境...
python --version >nul 2>&1
if %errorlevel% neq 0 (
    echo ❌ 错误: 未找到 Python。请确保已安装 Python 3.8 或更高版本并添加到 PATH 中。
    echo 💡 您可以从 https://www.python.org/downloads/ 下载 Python
    echo.
    pause
    exit /b 1
)

echo ✅ 找到 Python 版本:
python --version

REM 检查 NVIDIA 显卡和 CUDA
echo.
echo 🔍 检查 GPU 和 CUDA 环境...
nvidia-smi >nul 2>&1
if %errorlevel% equ 0 (
    echo ✅ 检测到 NVIDIA GPU:
    nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader,nounits
    echo.
    echo 🚀 将安装 CUDA 版本的 PyTorch 以启用 GPU 加速
    set "USE_CUDA=1"
) else (
    echo ⚠️  未检测到 NVIDIA GPU 或驱动，将使用 CPU 版本
    set "USE_CUDA=0"
)

echo.
echo 🔧 正在创建虚拟环境...
if not exist venv (
    python -m venv venv
    if %errorlevel% neq 0 (
        echo ❌ 错误: 创建虚拟环境失败
        pause
        exit /b 1
    )
    echo ✅ 虚拟环境创建成功
) else (
    echo ✅ 虚拟环境已存在
)

echo.
echo 🔄 正在激活虚拟环境...
call venv\Scripts\activate.bat

echo 📍 当前 Python 路径:
where python

echo.
echo 📦 正在升级 pip...
python -m pip install --upgrade pip

echo.
if "%USE_CUDA%"=="1" (
    echo 🚀 安装支持 GPU 加速的依赖包...
    echo 📥 正在安装 PyTorch CUDA 版本...
    python -m pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu118
    
    if %errorlevel% neq 0 (
        echo ⚠️  CUDA 版本安装失败，尝试安装 CPU 版本...
        python -m pip install torch torchaudio
    )
    
    echo 📥 正在安装其他 GPU 加速依赖...
    python -m pip install nvidia-ml-py3 accelerate
) else (
    echo 💻 安装 CPU 版本的依赖包...
    python -m pip install torch torchaudio
)

echo.
echo 📥 正在安装 Whisper 和其他依赖...
python -m pip install openai-whisper>=20231117 ffmpeg-python

if %errorlevel% neq 0 (
    echo ❌ 错误: 安装依赖包失败
    echo 💡 请检查网络连接或尝试使用国内镜像：
    echo    pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt
    pause
    exit /b 1
)

echo.
echo 🧪 正在测试安装...
python -c "import torch; print(f'PyTorch: {torch.__version__}')"
python -c "import torch; print(f'CUDA Available: {torch.cuda.is_available()}')"
if "%USE_CUDA%"=="1" (
    python -c "import torch; print(f'GPU Name: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"N/A\"}')"
)
python -c "import whisper; print('Whisper: OK')"

echo.
echo ================================
echo ✅ 安装完成！
echo ================================
echo.
if "%USE_CUDA%"=="1" (
    echo 🚀 GPU 加速已启用！您的 RTX 3080 将大幅提升转录速度
    echo ⚡ 预期性能提升: 5-10倍
) else (
    echo 💻 使用 CPU 模式运行
)
echo.
echo 📋 使用方法:
echo   批量处理: node main.js
echo   断点续传: node main.js --continue
echo   测试功能: node test_whisper.js
echo.
echo 💡 推荐的 Whisper 模型:
echo   - small:  平衡速度和质量（推荐）
echo   - base:   更快速度
echo   - medium: 更高质量（需要更多显存）
echo   - large:  最佳质量（需要8GB+ 显存）
echo.
pause